import random

import numpy as np

from scheduling_environment.jobShop import JobShop
from scheduling_environment.operation import Operation
from solution_methods.GA.src.heuristics import global_load_balancing_scheduler, local_load_balancing_scheduler, random_scheduler


def select_next_operation_from_job(jobShopEnv: JobShop, job_id) -> Operation:
    """
    从给定的作业车间环境中选择下一个可调度的操作。

    参数:
        jobShopEnv (JobShop): 作业车间环境。
        job_id: 作业ID。

    返回:
        Operation: 下一个可调度的操作。
    """
    # 遍历所有可调度的操作，找到与job_id匹配的操作并返回
    for operation in jobShopEnv.operations_available_for_scheduling:
        if operation.job_id == job_id:
            return operation


def pox_crossover(ind1, ind2, nr_preserving_jobs):
    """
    使用部分映射交叉（POX）对两个个体进行交叉操作。

    参数:
        ind1: 第一个个体。
        ind2: 第二个个体。
        nr_preserving_jobs: 保留的作业数量。

    返回:
        tuple: 两个新个体。
    """
    # 随机选择nr_preserving_jobs个作业作为保留作业
    preserving_jobs = random.sample(range(1, max(ind1)), nr_preserving_jobs)

    # 构建新的序列，保留部分来自ind1，其余来自ind2
    new_sequence_ind1 = list(filter(lambda a: a not in preserving_jobs, ind2))
    for i in range(len(ind1)):
        if ind1[i] in preserving_jobs:
            new_sequence_ind1.insert(i, ind1[i])

    # 构建新的序列，保留部分来自ind2，其余来自ind1
    new_sequence_ind2 = list(filter(lambda a: a not in preserving_jobs, ind1))
    for i in range(len(ind2)):
        if ind2[i] in preserving_jobs:
            new_sequence_ind2.insert(i, ind2[i])

    return new_sequence_ind1, new_sequence_ind2


def mutate_shortest_proc_time(individual, indpb, jobShopEnv: JobShop):
    """
    对个体进行变异操作，选择最短处理时间的机器。

    参数:
        individual: 个体。
        indpb: 每个基因变异的概率。
        jobShopEnv (JobShop): 作业车间环境。

    返回:
        individual: 变异后的个体。
    """
    # 对每个基因以indpb概率进行变异，选择最短处理时间的机器
    for i, _ in enumerate(individual):
        if random.random() < indpb:
            operation = jobShopEnv.operations[i]
            individual[i] = np.argmin(operation.processing_times)
    return individual


def mutate_sequence_exchange(individual, indpb):
    """
    对个体进行变异操作，交换两个随机位置的基因。

    参数:
        individual: 个体。
        indpb: 每个基因变异的概率。

    返回:
        individual: 变异后的个体。
    """
    # 对每个基因以indpb概率进行变异，随机交换两个位置的基因
    for i in range(len(individual)):
        if random.random() < indpb:
            j = random.choice([index for index in range(len(individual)) if index != i])
            individual[i], individual[j] = individual[j], individual[i]
    return individual


# 初始化遗传算法中的个体（使用随机动作选择启发式）
def init_individual(ind_class, jobShopEnv):
    """
    创建个体，个体是一个包含机器选择（选项索引）和操作序列（作业索引）的列表。

    参数:
        ind_class: 个体类。
        jobShopEnv (JobShop): 作业车间环境。

    返回:
        ind_class: 初始化后的个体。
    """
    rand = random.random()
    if rand <= 0.6:  # 60% 初始分配使用全局选择调度器
        jobShopEnv = global_load_balancing_scheduler(jobShopEnv)
    elif rand <= 0.9:  # 30% 初始分配使用局部选择调度器
        jobShopEnv = local_load_balancing_scheduler(jobShopEnv)
    else:  # 10% 初始分配使用随机调度器
        jobShopEnv = random_scheduler(jobShopEnv)

    # 获取操作序列和机器分配列表
    operation_sequence = [operation.job_id for operation in jobShopEnv.scheduled_operations]
    machine_selection = [
        (operation.operation_id, sorted(list(operation.processing_times.keys())).index(operation.scheduled_machine))
        for operation in jobShopEnv.scheduled_operations]
    machine_selection.sort()
    machine_selection = [allocation for _, allocation in machine_selection]
    jobShopEnv.reset()
    return ind_class([machine_selection, operation_sequence])


# 初始化种群
def init_population(toolbox, population_size):
    """
    初始化种群。

    参数:
        toolbox: DEAP工具箱。
        population_size: 种群大小。

    返回:
        list: 初始化后的种群。
    """
    return [toolbox.init_individual() for _ in range(population_size)]


def evaluate_individual(individual, jobShopEnv: JobShop, reset=True):
    """
    评估个体的适应度。

    参数:
        individual: 个体。
        jobShopEnv (JobShop): 作业车间环境。
        reset: 是否重置环境。

    返回:
        tuple: 适应度值和作业车间环境。
    """
    jobShopEnv.reset()
    jobShopEnv.update_operations_available_for_scheduling()
    for i in range(len(individual[0])):
        job_id = individual[1][i]
        operation = select_next_operation_from_job(jobShopEnv, job_id)
        operation_option_index = individual[0][operation.operation_id]
        machine_id = sorted(operation.processing_times.keys())[operation_option_index]
        duration = operation.processing_times[machine_id]

        jobShopEnv.schedule_operation_with_backfilling(operation, machine_id, duration)
        jobShopEnv.update_operations_available_for_scheduling()

    makespan = jobShopEnv.makespan

    if reset:
        jobShopEnv.reset()
    return makespan, jobShopEnv


def evaluate_population(toolbox, population):
    """
    评估种群中所有个体的适应度。

    参数:
        toolbox: DEAP工具箱。
        population: 种群。

    返回:
        list: 所有个体的适应度值。
    """
    # 并行评估种群
    population = [[ind[0], ind[1]] for ind in population]
    fitnesses = toolbox.map(toolbox.evaluate_individual, population)
    fitnesses = [(fit[0],) for fit in fitnesses]

    return fitnesses


def variation(population, toolbox, pop_size, cr, indpb):
    """
    对种群进行变异操作，包括交叉、复制和变异。

    参数:
        population: 种群。
        toolbox: DEAP工具箱。
        pop_size: 新种群大小。
        cr: 交叉概率。
        indpb: 基因变异概率。

    返回:
        list: 变异后的种群。
    """
    offspring = []
    for _ in range(int(pop_size)):
        op_choice = random.random()
        if op_choice < cr:  # 应用交叉
            ind1, ind2 = list(map(toolbox.clone, random.sample(population, 2)))
            if random.random() < 0.5:
                ind1[0], ind2[0] = toolbox.mate_TwoPoint(ind1[0], ind2[0])
            else:
                ind1[0], ind2[0] = toolbox.mate_Uniform(ind1[0], ind2[0])

            ind1[1], ind2[1] = toolbox.mate_POX(ind1[1], ind2[1])
            del ind1.fitness.values, ind2.fitness.values

        else:  # 应用复制
            ind1 = toolbox.clone(random.choice(population))

        # 应用变异
        ind1[0] = toolbox.mutate_machine_selection(ind1[0], indpb)
        ind1[1] = toolbox.mutate_operation_sequence(ind1[1], indpb)

        del ind1.fitness.values
        offspring.append(ind1)

    return offspring


def repair_precedence_constraints(env, offspring):
    """
    修复后代个体中的优先约束关系。

    参数:
        env (JobShop): 作业车间环境。
        offspring: 后代个体列表。

    返回:
        list: 修复后的后代个体列表。
    """
    precedence_relations = env.precedence_relations_jobs
    for ind in offspring:
        i = 0
        lst = ind[1]
        while i < len(ind[1]):
            if lst[i] in precedence_relations.keys():
                max_index = 0
                for j in precedence_relations[lst[i]]:
                    index = len(lst) - 1 - lst[::-1].index(j)
                    if index > max_index:
                        max_index = index
                if max_index > i:
                    item = lst[i]
                    lst.pop(i)  # 从源索引移除项目
                    lst.insert(max_index, item)
                    continue
            i += 1
    return offspring
